// PhotonMapping.cpp : Defines the entry point for the console application.
//

#include "stdafx.h"
#include "Scene.h"
#include "RayTracer.h"
#include "PhotonMap.h"
#include "getopt.h"

bool WriteTGA(const tstring& p_FN, const int p_Width, const int p_Height, const std::vector<PhotonMapping::Color>& p_Buffer)
{
	std::ofstream imageFile(p_FN.c_str(), std::ios_base::binary);
	if (!imageFile)
		return false;
	// specific TGA handling code
	imageFile.put(0).put(0);
	imageFile.put(2);                  /* uncompressed RGB */

	imageFile.put(0).put(0);
	imageFile.put(0).put(0);
	imageFile.put(0);

	imageFile.put(0).put(0);           /* X origin */
	imageFile.put(0).put(0);           /* y origin */

	const int widht = p_Width;
	const int height = p_Height;
	imageFile.put((widht & 0x00FF)).put((widht & 0xFF00) / 256);
	imageFile.put((height & 0x00FF)).put((height & 0xFF00) / 256);
	imageFile.put(24);                 /* 24 bit bitmap */
	imageFile.put(0);

	for(size_t i = 0; i < p_Buffer.size(); ++i)
	{
		const PhotonMapping::Color& curCol = p_Buffer[i];
		imageFile.put(static_cast<char>(Math::Min(curCol.m_B * 255.0f, 255.0f)));
		imageFile.put(static_cast<char>(Math::Min(curCol.m_G * 255.0f, 255.0f)));
		imageFile.put(static_cast<char>(Math::Min(curCol.m_R * 255.0f, 255.0f)));
	}
	return true;
}

struct TechniqueMap
{
	const tchar* m_Name;
	const uintx m_Value;
};

intx GetTechnique(const tstring& p_TechniqueName)
{
	using namespace PhotonMapping;

	const uintx MaxTechniques = 4;
	const TechniqueMap techniques[MaxTechniques] = 
	{
		{ _T("swkdtree"), PhotonMap::SWKdTree },
		{ _T("gpukdtree"), PhotonMap::GPUKdTree },
		{ _T("gpuhash"), PhotonMap::GPUHash },
		{ _T("gpubf"), PhotonMap::GPUBF }
	};

	for(uintx i = 0; i < MaxTechniques; ++i)
	{
		if(_tcsicmp(techniques[i].m_Name, p_TechniqueName.c_str()) == 0)
			return techniques[i].m_Value;
	}
	return -1;
}

int _tmain(int argc, _TCHAR* argv[])
{
 	using namespace PhotonMapping;
	
	// program parameters
	bool useGlobalPM = false;
	bool useCausticPM = false;
	bool generateGlobalPM = false;
	bool generateCausticPM = false;
	bool justRayTrace = false;
	tstring globalPMFile;
	tstring causticPMFile;
	uintx gMaxPhotons, cMaxPhotons;
	float gMaxRadius, cMaxRadius;
	uintx genMaxPhotons;
	tstring techniqueName;
	bool verboseMode = false;
	int width = 640, height = 480;

	// parse arguments
	intx c;
	while((c = getopt(argc, argv, _T("vrw:h:x:g:c:t:"))) != -1)
	{
		switch (c)
		{
			case _T('r'):
				justRayTrace = true;
				break;
			case _T('v'):
				verboseMode = true;
				break;
			case _T('h'):
				height = _tstoi(optarg);
				break;
			case _T('w'):
				width = _tstoi(optarg);
				break;
			case _T('x'):
				// generate global pm
				generateCausticPM = optarg[0] == _T('c');
				generateGlobalPM = !generateCausticPM;
				genMaxPhotons = _tstoi(argv[optind]);
				optind += 1;
				break;
			case _T('g'):
				// use global photon map (with settings)
				useGlobalPM = true;
				globalPMFile = optarg;
				gMaxPhotons = _tstoi(argv[optind]);
				gMaxRadius = static_cast<float>(_tstof(argv[optind + 1]));
				optind += 2;
				break;
			case _T('c'):
				// use caustic photon map (with settings)
				useCausticPM = true;
				causticPMFile = optarg;
				cMaxPhotons = _tstoi(argv[optind]);
				cMaxRadius = static_cast<float>(_tstof(argv[optind + 1]));
				optind += 2;
				break;
			case _T('t'):
				techniqueName = optarg;
				break;
			case _T('?'):
				_tprintf(_T("Usage: cachesim [-g <name max r>|-c <name max r>] -t <tech> [outputfile]\n\n"));
				_tprintf(_T("  g: Use global photon map (each query uses <max> photons and a query radius of <r>)\n"));
				_tprintf(_T("  c: Use caustic photon map (each query uses <max> photons and a query radius of <r>)\n"));
				_tprintf(_T("  t: Photon map technique: swkdtree, gpukdtree, gpuhash or gpubf\n"));
				_tprintf(_T("  x: Generate photon map (c for caustic or g for global, max photons)\n"));
				_tprintf(_T("  w: Image width size in pixels (default %d)\n"), width);
				_tprintf(_T("  h: Image height size in pixels (default %d)\n"), height);
				_tprintf(_T("  r: Render image using pure raytracing\n"));
				_tprintf(_T("  v: Verbose mode - won't produce any image outputs\n"));
				return -1;
			default:
				break;
		}
	}

	// generate output filename
	tstring outputName;
	if(optind >= argc)
	{
		if(generateCausticPM || generateGlobalPM || justRayTrace)
		{
			_tprintf(_T("No output file specified!\n"));
			return -1;
		}

		tostringstream stream;
		stream << _T("result");
		if(useGlobalPM)
			stream << _T("_") << globalPMFile << _T("_") << gMaxPhotons << _T("_") << gMaxRadius;
		if(useCausticPM)
			stream << _T("_") << causticPMFile << _T("_") << cMaxPhotons << _T("_") << cMaxRadius;
		stream << _T("_") << techniqueName << _T("_") << width << _T("_") << height;
		outputName = stream.str();	
	}
	else
		outputName = argv[optind];

	// load scene
	ScenePtr scene = Scene::Load(_T("scene.txt"));
	if(!scene)
		return -1;

	// create photon map
	if(generateCausticPM || generateGlobalPM)
	{
		_tprintf(_T("Creating photon map ...\n"));
		PhotonMapPtr pm = PhotonMapPtr(new PhotonMap(0.0f, 0, 0));
		pm->Create(scene, genMaxPhotons, generateCausticPM);
		PhotonMap::Store(outputName.c_str(), pm);
		return 0;
	}

	if(justRayTrace)
	{
		_tprintf(_T("Ray tracing ...\n"));
		std::vector<Color> buffer;
		RayTracer rt;
		rt.Draw(scene, buffer, width, height);
		const tstring resultName = outputName + _T(".tga");
		WriteTGA(resultName, width, height, buffer);
		return 0;
	}

	// determine technique
	intx technique = GetTechnique(techniqueName);
	if(technique < 0)
	{
		_tprintf(_T("Unknown technique specified: %s!\n"), techniqueName);
		return -1;
	}

	// redirect stdout to a log file
	const tstring logFileName = outputName + _T(".log");
	_tfreopen(logFileName.c_str(), _T("w"), stdout);

	// print parameter information
	_tprintf(_T("Running with parameters:\n"));
	if(useGlobalPM)
		_tprintf(_T("  Global Photon Map (%s\tMaxPhotons: %d\tMaxRadius: %.2f)\n"), globalPMFile.c_str(), gMaxPhotons, gMaxRadius);
	if(useCausticPM)
		_tprintf(_T("  Caustic Photon Map (%s\tMaxPhotons: %d\tMaxRadius: %.2f)\n"), causticPMFile.c_str(), cMaxPhotons, cMaxRadius);
	_tprintf(_T("  Output: %d x %d\n"), width, height);
	_tprintf(_T("  Using technique %s\n"), techniqueName.c_str());

	// CUDA Initialisation
	cuInit(0);

	// because the first cudaMalloc call usually takes a long time (> 1 sec.) we 
	// make a dummy allocation here to avoid having that delay in our performance
	// measurements
	int* dummy;
	cutilSafeCall(cudaMalloc((void**)&dummy, sizeof(int)));
	cutilSafeCall(cudaFree(dummy));

	// load photon maps
	PhotonMapPtr globalPM, causticPM;
	if(useGlobalPM)
	{
		_tprintf(_T("Loading global photon map ...\n"));
		globalPM = PhotonMap::Load(globalPMFile.c_str(), gMaxRadius, gMaxPhotons, technique);
	}
	if(useCausticPM)
	{
		_tprintf(_T("Loading caustic photon map ...\n"));
		causticPM = PhotonMap::Load(causticPMFile.c_str(), cMaxRadius, cMaxPhotons, technique);
	}

	// prepare for rendering
	_tprintf(_T("Preparing photon maps ...\n"));
	if(globalPM) globalPM->Prepare();
	if(causticPM) causticPM->Prepare();

	// render using our photon maps
	_tprintf(_T("Ray tracing photon map ...\n"));
	std::vector<Color> buffer;
	std::vector<Color> bufferGlobalPM;
	std::vector<Color> bufferCausticPM;
	RayTracer rt;

	rt.Draw(scene, globalPM, causticPM, buffer, width, height);
 	if(globalPM)
		rt.DrawPhotonMap(scene, globalPM, width, height, bufferGlobalPM);
	if(causticPM)
		rt.DrawPhotonMap(scene, causticPM, width, height, bufferCausticPM);

	// write results back
	if(!verboseMode)
	{
		const tstring resultName = outputName + _T(".tga");
		WriteTGA(resultName, width, height, buffer);
		if(useGlobalPM)
		{
			const tstring resultGlobalPMName = outputName + _T("_globalpm.tga");
			WriteTGA(resultGlobalPMName, width, height, bufferGlobalPM);
		}
		if(useCausticPM)
		{
			const tstring resultCausticPMName = outputName + _T("_causticpm.tga");
			WriteTGA(resultCausticPMName, width, height, bufferCausticPM);
		}
	}

	cudaThreadExit();
	return 0;
}

